import numpy as np
import pandas as pd
from python_ai.common.xcommon import sep
pd.set_option('display.max_columns', None)

df = pd.read_csv('bayes_xinxi.txt')

y = df.iloc[:, 0]
x = df.iloc[:, 1]

from sklearn.feature_extraction.text import CountVectorizer
tf_model = CountVectorizer()
x = tf_model.fit_transform(x)

from sklearn.naive_bayes import MultinomialNB
model = MultinomialNB()
model.fit(x, y)
print(f'Training score: {model.score(x, y)}')

X_ = ['Chinese Chinese Chinese Tokyo Japan']
X_ = tf_model.transform(X_)
h = model.predict(X_)
print(f'Predict: {h}')
